package ca.compsci.opent.compiler.codegen;

import static org.objectweb.asm.Opcodes.*;

import java.util.Iterator;

import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import ca.compsci.opent.compiler.analysis.*;
import ca.compsci.opent.compiler.node.*;
import ca.compsci.opent.compiler.semantics.*;

public class BytecodeInsnGenerator extends DepthFirstAdapter {
	private MethodVisitor mv;
	
	private SymbolTable scope;
	private Iterator<SymbolTable> subscopes;
	private int max_stack_size = 0, max_locals_size = 0;
	private Label lbl_exit_loop = null, lbl_exit_if = null;
	
	private Type last_type = null;
	
	public BytecodeInsnGenerator(MethodVisitor mv, int locals_size, SymbolTable scope) {
		this.mv = mv;
		this.scope = scope;
		this.subscopes = scope.getSubScopeIterator();
		updateLocalsSize(locals_size);
	}
	
	@Override
	public void caseABlockStmt(ABlockStmt node) {
		SymbolTable current_scope = scope;
		Iterator<SymbolTable> current_sub_scopes = subscopes;
		Label current_lbl_exit_if_else = lbl_exit_if;
		
		lbl_exit_if = null;
		scope = subscopes.next();
		subscopes = scope.getSubScopeIterator();
		super.caseABlockStmt(node);
		subscopes = current_sub_scopes;
		scope = current_scope;
		lbl_exit_if = current_lbl_exit_if_else;
	}

	@Override
	public void inAPutStmt(APutStmt node) {
		mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
	}
	
	@Override
	public void outAPutStmt(APutStmt node) {
		switch (last_type) {
			case Real64:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(D)V");
				break;
			case Real32:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(F)V");
				break;
			case Integer64:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(J)V");
				break;
			case Integer32:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(I)V");
				break;
			case Integer8:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(B)V");
				break;
			case Boolean:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Z)V");
				break;
			case String:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V");
				break;
			default: // Possible error handling should go here
				break;
		}
		
		if (is64Bit(last_type)) // account for the `out`
			updateStackSize(5); // 4 + 1
		else
			updateStackSize(3); // 2 + 1
	}

	@Override
	public void outAInferDeclStmt(AInferDeclStmt node) {
		TIdentifier first = node.getIdentifier();
		Variable var = (Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(first));
		allocateVariable(var);
		for (TIdentifier ident : node.getRest())
			allocateVariable((Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(ident)));
	}
	
	@Override
	public void outAInferInitStmt(AInferInitStmt node) {
		TIdentifier first = node.getIdentifier();
		Variable var = (Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(first));
		allocateVariable(var);
		dupTopValue(var.getType());
		storePushed(var);  // POPs
		for (TIdentifier ident : node.getRest()) {
			var = (Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(ident));
			allocateVariable(var);
			dupTopValue(var.getType());
			storePushed(var); // POPs
		}
		popTopValue(var.getType());
	}		

	@Override
	public void outATypedDeclStmt(ATypedDeclStmt node) {
		TIdentifier first = node.getIdentifier();
		allocateVariable((Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(first)));
		for (TIdentifier ident : node.getRest())
			allocateVariable((Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(ident)));
	}
	
	@Override
	public void outATypedInitStmt(ATypedInitStmt node) {
		TIdentifier first = node.getIdentifier();
		Variable var = (Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(first));
		allocateVariable(var);
		dupTopValue(var.getType());
		storePushed(var);  // POPs
		for (TIdentifier ident : node.getRest()) {
			var = (Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(ident));
			allocateVariable(var);
			dupTopValue(var.getType());
			storePushed(var); // POPs
		}
		popTopValue(var.getType());
	}
	
	@Override
	public void caseALoopStmt(ALoopStmt node) {
		Label current_exit_loop_lbl = lbl_exit_loop;
		Label lbl_loop = new Label();
		
		lbl_exit_loop = new Label();
		
		mv.visitLabel(lbl_loop);
		super.caseALoopStmt(node);
		mv.visitJumpInsn(GOTO, lbl_loop);
		mv.visitLabel(lbl_exit_loop);
		
		lbl_exit_loop = current_exit_loop_lbl;
	}
	
	@Override
	public void caseAIfStmt(AIfStmt node) {
		Label lbl_if_false = new Label();
		boolean created_lbl_exit_if = false;
		
		if (lbl_exit_if == null) {
			lbl_exit_if = new Label();
			created_lbl_exit_if = true;
		}
		
		PExpr expr = node.getExpr();
		PStmt if_body = node.getIfBody();
		PStmt else_body = node.getElseBody();
		
		if (expr != null)
			expr.apply(this);
		
		mv.visitJumpInsn(IFEQ, lbl_if_false); // if top value is FALSE	
		if (if_body != null)
			if_body.apply(this);
		mv.visitJumpInsn(GOTO, lbl_exit_if);
		
		mv.visitLabel(lbl_if_false);
		for (PStmt stmt : node.getElsifs())
			stmt.apply(this);
		
		if (else_body != null)
			else_body.apply(this);
		
		if (created_lbl_exit_if) {
			mv.visitLabel(lbl_exit_if);
			lbl_exit_if = null;
		}
	}
	
	@Override
	public void outAExitWhenStmt(AExitWhenStmt node) {
		mv.visitJumpInsn(IFNE, lbl_exit_loop); // if top value is TRUE
	}

	@Override
	public void outAExprStmt(AExprStmt node) {
		// expr stmts push 1 value onto stack, so POP it
		popTopValue(last_type);
	}
	
	@Override
	public void outAAssignExpr(AAssignExpr node) {
		TIdentifier ident = node.getIdentifier();
		Variable var = (Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(ident));
		dupTopValue(var.getType());
		storePushed(var); // POPs
		updateStackSize(2);
	}
	
	@Override
	public void caseABinOprExpr(ABinOprExpr node) {
		final MethodVisitor mv = this.mv;
		final Label lbl_success, lbl_next;
		
		POpr opr = node.getOpr();
		
		Expression expr = (Expression) scope.getInfo(SemanticAnalyzer.mkABinOprExprKey(node));
		
		final Type promoted_type = expr.getPromotedType();
		
		if (opr instanceof AAndOpr || opr instanceof AOrOpr) {
			lbl_success = new Label();
			lbl_next = new Label();
		} else {
			lbl_success = null;
			lbl_next = null;
		}
		
		if (node.getLeft() != null) {
			node.getLeft().apply(this);
			if (opr instanceof AOrOpr) {
				mv.visitJumpInsn(IFNE, lbl_success); // if top value on stack is TRUE
			} else if (opr instanceof AAndOpr) {
				mv.visitJumpInsn(IFEQ, lbl_success); // if top value on stack is FALSE
			} else if (opr instanceof APowerOpr) {
				handleTypeConversion(last_type, Type.Real64);
		    }  else {
				handleTypeConversion(last_type, expr.getPromotedType());
			}
		}
		
		if (node.getRight() != null) {
			node.getRight().apply(this);
			if (opr instanceof AOrOpr) {
				mv.visitJumpInsn(IFNE, lbl_success); // if top value on stack is TRUE
			} else if (opr instanceof AAndOpr) {
				mv.visitJumpInsn(IFEQ, lbl_success); // if top value on stack is FALSE
			} else if (opr instanceof APowerOpr) {
				handleTypeConversion(last_type, Type.Real64);
		    } else {
				handleTypeConversion(last_type, expr.getPromotedType());
			}
		}
		
		if (opr != null) {
			opr.apply(new AnalysisAdapter() {
				@Override
				public void caseAAddOpr(AAddOpr node) {
					if (promoted_type.equals(Type.String))
						mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/String", "concat", "(Ljava/lang/String;)Ljava/lang/String;");
					else if (is64Bit(promoted_type))
						mv.visitInsn(LADD);
					else
						mv.visitInsn(IADD);
				}

				@Override
				public void caseASubOpr(ASubOpr node) {
					if (is64Bit(promoted_type))
						mv.visitInsn(LSUB);
					else
						mv.visitInsn(ISUB);
				}

				@Override
				public void caseADivOpr(ADivOpr node) {
					if (is64Bit(promoted_type))
						mv.visitInsn(LDIV);
					else
						mv.visitInsn(IDIV);
				}

				@Override
				public void caseAMulOpr(AMulOpr node) {
					if (is64Bit(promoted_type))
						mv.visitInsn(LMUL);
					else
						mv.visitInsn(IMUL);
				}

				@Override
				public void caseAShlOpr(AShlOpr node) {
					if (is64Bit(promoted_type))
						mv.visitInsn(LSHL);
					else
						mv.visitInsn(ISHL);
				}

				@Override
				public void caseAShrOpr(AShrOpr node) {
					if (is64Bit(promoted_type))
						mv.visitInsn(LSHR);
					else
						mv.visitInsn(ISHR);
				}

				@Override
				public void caseAEqOpr(AEqOpr node) {
					if (promoted_type.equals(Type.String)) {
						mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/String", "compareTo", "(Ljava/lang/String;)I");
						handleBooleanInsn(IFEQ);
					} else {
						handleBooleanInsn(IF_ICMPEQ);
					}
				}
				
				@Override
				public void caseANoteqOpr(ANoteqOpr node) {
					if (promoted_type.equals(Type.String)) {
						mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/String", "compareTo", "(Ljava/lang/String;)I");
						handleBooleanInsn(IFNE);
					} else {
						handleBooleanInsn(IF_ICMPNE);
					}
				}

				@Override
				public void caseAGteqOpr(AGteqOpr node) {
					if (promoted_type.equals(Type.String)) {
						mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/String", "compareTo", "(Ljava/lang/String;)I");
						handleBooleanInsn(IFGE);
					} else {
						handleBooleanInsn(IF_ICMPGE);
					}
				}

				@Override
				public void caseAGtOpr(AGtOpr node) {
					if (promoted_type.equals(Type.String)) {
						mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/String", "compareTo", "(Ljava/lang/String;)I");
						handleBooleanInsn(IFGT);
					} else {
						handleBooleanInsn(IF_ICMPGT);
					}
				}

				@Override
				public void caseALteqOpr(ALteqOpr node) {
					if (promoted_type.equals(Type.String)) {
						mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/String", "compareTo", "(Ljava/lang/String;)I");
						handleBooleanInsn(IFLE);
					} else {
						handleBooleanInsn(IF_ICMPLE);
					}
				}

				@Override
				public void caseALtOpr(ALtOpr node) {
					if (promoted_type.equals(Type.String)) {
						mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/String", "compareTo", "(Ljava/lang/String;)I");
						handleBooleanInsn(IFLT);
					} else {
						handleBooleanInsn(IF_ICMPLT);
					}
				}
				
				@Override
				public void caseAModOpr(AModOpr node) {
					if (is64Bit(promoted_type))
						mv.visitInsn(LREM);
					else
						mv.visitInsn(IREM);
				}

				@Override
				public void caseAOrOpr(AOrOpr node) {
					mv.visitInsn(ICONST_0);				 // PUSH FALSE
					mv.visitJumpInsn(GOTO, lbl_next);	 // GOTO Next
					mv.visitLabel(lbl_success);			 // Success...
					mv.visitInsn(ICONST_1);				 // PUSH TRUE
					mv.visitLabel(lbl_next);			 // Next...
				}

				@Override
				public void caseAAndOpr(AAndOpr node) {
					mv.visitInsn(ICONST_1);				 // PUSH TRUE
					mv.visitJumpInsn(GOTO, lbl_next);    // GOTO Next
					mv.visitLabel(lbl_success);          // Success...
					mv.visitInsn(ICONST_0);				 // PUSH FALSE
					mv.visitLabel(lbl_next);		     // Next..
				}

				@Override
				public void caseAPowerOpr(APowerOpr node) {
					mv.visitMethodInsn(INVOKESTATIC, "java/lang/Math", "pow", "(DD)D");
					handleTypeConversion(Type.Real64, promoted_type);
				}
			});
			last_type = expr.getType();
		}
		
		if (is64Bit(promoted_type))
			updateStackSize(4);
		else
			updateStackSize(2);
	}

	@Override
	public void outAUnrOprExpr(AUnrOprExpr node) {
		node.getOpr().apply(new AnalysisAdapter() {

			@Override
			public void caseAIdentityOpr(AIdentityOpr node) {
				// ignore... (identity after all)
			}

			@Override
			public void caseANegateOpr(ANegateOpr node) {
				if (is64Bit(last_type))
					mv.visitInsn(LNEG);
				else
					mv.visitInsn(INEG);
			}

			@Override
			public void caseANotOpr(ANotOpr node) {
				Label lbl_iftrue = new Label(), lbl_next = new Label();
				mv.visitJumpInsn(IFEQ, lbl_iftrue);
				mv.visitInsn(ICONST_0);		// PUSH FALSE
				mv.visitJumpInsn(GOTO, lbl_next);
				mv.visitLabel(lbl_iftrue);
				mv.visitInsn(ICONST_1);		// PUSH TRUE
				mv.visitLabel(lbl_next);
			}
		});
	}

	@Override
	public void outAIdentifierExpr(AIdentifierExpr node) {
		Variable var = (Variable) scope.getInfo(SemanticAnalyzer.mkTIdentifierKey(node.getIdentifier()));
		loadVariable(var);
		last_type = var.getType();
	}
	
	@Override
	public void outANumberExpr(ANumberExpr node) {
		Literal num_lit = (Literal) scope.getInfo(SemanticAnalyzer.mkANumberExprKey(node));
		Number num = (Number) num_lit.getValue();
		switch (num_lit.getType()) {
		case Integer64:
			mv.visitLdcInsn(num.longValue());
			break;
		case Integer32:	
			mv.visitLdcInsn(num.intValue());
			break;
		case Integer16:
			mv.visitLdcInsn(num.shortValue());
			break;
		case Integer8:
			mv.visitLdcInsn(num.byteValue());
			break;
		case Real32:
			mv.visitLdcInsn(num.floatValue());
			break;
		case Real64:
			mv.visitLdcInsn(num.doubleValue());
			break;
		}
		updateStackSize(1);
		last_type = num_lit.getType();
	}
	
	@Override
	public void outABooleanExpr(ABooleanExpr node) {
		Literal bool = (Literal) scope.getInfo(SemanticAnalyzer.mkABooleanExprKey(node));
		if ((Boolean) bool.getValue())
			mv.visitInsn(ICONST_1);  // TRUE
		else
			mv.visitInsn(ICONST_0);  // FALSE
		last_type = bool.getType();
	}
	
	@Override
	public void outAStringExpr(AStringExpr node) {
		Literal str = (Literal) scope.getInfo(SemanticAnalyzer.mkAStringExprKey(node));
		mv.visitLdcInsn((String)str.getValue());
		updateStackSize(1);
		last_type = str.getType();
	}
	
	public int getStackSize() {
		return max_stack_size;
	}
	
	public int getLocalsSize() {
		return max_locals_size;
	}
	
	// Private Helper methods
	
	private void updateStackSize(int size) {
		if (size > max_stack_size)
			max_stack_size = size;
	}
	
	private void updateLocalsSize(int size) {
		if (size > max_locals_size)
			max_locals_size = size;
	}
	
	private void allocateVariable(Variable var) {
		var.setId(max_locals_size);
		updateLocalsSize(max_locals_size + 1);
	}
	
	// pops 2 and push 1 to stack
	private void handleBooleanInsn(int jump_insn) {
		Label lbl_iftrue = new Label();
		Label lbl_next = new Label();
		
		mv.visitJumpInsn(jump_insn, lbl_iftrue); // jump if true
		mv.visitInsn(ICONST_0); // FALSE
		mv.visitJumpInsn(GOTO, lbl_next);
		
		mv.visitLabel(lbl_iftrue); // establish lbl_iftrue jump point
		mv.visitInsn(ICONST_1);  // TRUE
		
		mv.visitLabel(lbl_next); // establish lbl_next jump point
	}
	
	// pops 1 values from stack
	private void storePushed(Variable var) {
		int id = var.getId();
		switch (var.getType()) {
		case Integer8:
		case Integer16:
		case Integer32:
		case Boolean:
			mv.visitVarInsn(ISTORE, id);
			break;
		case Integer64:
			mv.visitVarInsn(LSTORE, id);
			break;
		case Real32:
			mv.visitVarInsn(FSTORE, id);
			break;
		case Real64:
			mv.visitVarInsn(DSTORE, id);
			break;
		default:
			mv.visitVarInsn(ASTORE, id);
			break;
		}
		updateLocalsSize(id + 1); // id's start from 0
	}
	
	// pushes 1 value to stack
	private void loadVariable(Variable var) {
		int id = var.getId();
		switch (var.getType()) {
		case Integer8:
		case Integer16:
		case Integer32:
		case Boolean:
			mv.visitVarInsn(ILOAD, id);
			break;
		case Integer64:
			mv.visitVarInsn(LLOAD, id);
			break;
		case Real32:
			mv.visitVarInsn(FLOAD, id);
			break;
		case Real64:
			mv.visitVarInsn(DLOAD, id);
			break;
		default:
			mv.visitVarInsn(ALOAD, id);
			break;
		}
		updateStackSize(1);
	}
	
	private void popTopValue(Type t) {
		if (is64Bit(t))
			mv.visitInsn(POP2);
		else
			mv.visitInsn(POP);
	}
	
	private void dupTopValue(Type t) {
		if (is64Bit(t))
			mv.visitInsn(DUP2);
		else
			mv.visitInsn(DUP);
	}
	
	private void handleTypeConversion(Type from_type, Type to_type) {
		if (!from_type.equals(to_type)) {
			switch (from_type) {
			case Real64:
				switch (to_type) {
				case Real32:
					mv.visitInsn(D2F);
					break;
				case Integer8:
				case Integer16:
				case Integer32:
					mv.visitInsn(D2I);
					break;
				case Integer64:
					mv.visitInsn(D2L);
					break;
				case String:
					mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(D)Ljava/lang/String;");
					break;
				}
				break;
			case Real32:
				switch (to_type) {
				case Real64:
					mv.visitInsn(F2D);
					break;
				case Integer8:
				case Integer16:
				case Integer32:
					mv.visitInsn(F2I);
					break;
				case Integer64:
					mv.visitInsn(F2L);
					break;
				case String:
					mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(F)Ljava/lang/String;");
					break;
				}
				break;
			case Integer8:
			case Integer16:
			case Integer32:
				switch (to_type) {
				case Real64:
					mv.visitInsn(I2D);
					break;
				case Real32:
					mv.visitInsn(I2F);
					break;
				case Integer64:
					mv.visitInsn(I2L);
					break;
				case String:
					switch (from_type) {
					case Integer16:
						mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(S)Ljava/lang/String;");
						break;
					case Integer32:
						mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(I)Ljava/lang/String;");
						break;
					case Integer8:
						mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(B)Ljava/lang/String;");
						break;
					}
					break;
				}
				break;
			case Integer64:
				switch (to_type) {
				case Real64:
					mv.visitInsn(L2D);
					break;
				case Real32:
					mv.visitInsn(L2F);
					break;
				case Integer8:
				case Integer16:
				case Integer32:
					mv.visitInsn(L2I);
					break;
				case String:
					mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(J)Ljava/lang/String;");
					break;
				}
				break;
			case Boolean:
				case String:
					mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(Z)Ljava/lang/String;");
					break;
			}
		}
	}
	
	private boolean is64Bit(Type type) {
		switch (type) {
			case Integer64:
			case Real64:
				return true;
			default:
				return false;
		}
	}
}
